#include "proxy_kpabe.h"
#include "kpabe_KPAbe.h"
#include "string.h"
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
using namespace std;

extern kpabe_pub_t* pub;
unordered_map<int, kpabe_pub_t*> streamMap;
unordered_map<int, kpabe_prv_t*> policyMap;

kpabe_cph_t* convertToCipher(kpabe_pub_t* pubK, unsigned char* cipher, int cipherLen){
	GByteArray *buf =  convertToByteArray(cipher, cipherLen);
	return kpabe_cph_unserialize(pubK, buf, 1);
}

int initACInstance(int streamId, unsigned char* pubBytes, int pubLen, int policyId, unsigned
char* prvBytes, int prvLen){
	pub = convertToPubKey(pubBytes, pubLen);	
	kpabe_prv_t* prv = convertToPrvKey(prvBytes, prvLen);

	unordered_map<int, kpabe_pub_t*>::const_iterator it = streamMap.find(streamId);
	if (it==streamMap.end())
{	
		streamMap[streamId] = pub;
}	
	policyMap[policyId] = prv;

	return 0;	
}
JNIEXPORT void JNICALL Java_pCloud_crypto_kpabe_KPAbe_initACInstance
  (JNIEnv *env, jobject obj, jbyteArray pubB, jint streamId, jbyteArray tkB, jint policyId){
	int pubL = env->GetArrayLength(pubB);
	unsigned char* pubBytes = (unsigned char*)malloc(pubL*sizeof(char));
	env->GetByteArrayRegion(pubB,0,pubL,(jbyte*)pubBytes);

	int prvL = env->GetArrayLength(tkB);
	unsigned char* prvBytes = (unsigned char*)malloc(prvL*sizeof(char));
	env->GetByteArrayRegion(tkB,0,prvL,(jbyte*)prvBytes);
	
	initACInstance(streamId, pubBytes, pubL, policyId, prvBytes, prvL);		
}

//transform the ciphertext, using prv key indexed by the policyId
int transform_simple(unsigned char* cipher, int cipherLen, int streamId, int policyId, unsigned char** result,
int *len){
	kpabe_pub_t* pub = streamMap[streamId];
	kpabe_prv_t* prv = policyMap[policyId];
	kpabe_cph_t* cph = convertToCipher(pub, cipher, cipherLen); 		

	element_t userCipher;
	int s = kpabe_transform(pub,prv,cph,userCipher);
	*len = element_length_in_bytes(userCipher);
	*result = (unsigned char*) malloc((*len)*sizeof(char));
	element_to_bytes(*result, userCipher);
	return s;
}

JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_transform_1simple
  (JNIEnv *env, jobject obj, jbyteArray cipher, jint streamId, jint policyId){
	int cL = env->GetArrayLength(cipher);
	unsigned char* cipherBytes = (unsigned char*)malloc(cL*sizeof(char));
	env->GetByteArrayRegion(cipher,0,cL,(jbyte*)cipherBytes);
	
	unsigned char* transformedCipher;
	int transformedCipherLen;
	transform_simple(cipherBytes, cL, streamId, policyId, &transformedCipher,
&transformedCipherLen);
	jbyteArray ret = env->NewByteArray(transformedCipherLen);
	env->SetByteArrayRegion(ret,0,transformedCipherLen,(jbyte*)transformedCipher);
	return ret;
}

//multiplying all ciphers
int multiplyCiphers(unsigned char** ciphers, int size, int streamId, unsigned char** result, int*
len){
	kpabe_pub_t* pub = streamMap[streamId];
	element_t one, prod, cipher;
	element_init_GT(one, pub->p);
	element_set_si(one,1);
	element_init_GT(prod, pub->p);
	element_init_GT(cipher, pub->p);
	for (int i=0; i<size; i++){
		element_from_bytes(cipher, ciphers[i]);
		element_mul(prod,cipher,one);
		element_set(one,prod);
	}	
	*len = element_length_in_bytes(prod);
	*result = (unsigned char*) malloc((*len)*sizeof(char));
	element_to_bytes(*result,prod);
	return 0;
}
JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_multiplyCiphers
  (JNIEnv *env, jobject obj, jobjectArray ciphers, jint streamId){
	int size = env->GetArrayLength(ciphers);
	unsigned char** cipherArray = (unsigned char**)malloc(size*sizeof(char*));
	for (int i=0; i<size; i++){
		jbyteArray tmp = (jbyteArray)env->GetObjectArrayElement(ciphers,i);
		int len = env->GetArrayLength(tmp);
		cipherArray[i] = (unsigned char*)malloc(len*sizeof(char));
		env->GetByteArrayRegion(tmp,0,len,(jbyte*)cipherArray[i]);
	}
	
	unsigned char *result;
	int resultLen;
	multiplyCiphers(cipherArray,size,streamId,&result,&resultLen);
	
	jbyteArray ret = env->NewByteArray(resultLen);
	env->SetByteArrayRegion(ret,0,resultLen,(jbyte*)result);
	return ret;
}

